/////////////////////////////////////////////////////////////////////////////
// Original code from libhdfs3. Copyright (c) 2013 - 2014, Pivotal Inc.
// All rights reserved. Author: Zhanwei Wang
/////////////////////////////////////////////////////////////////////////////
//  Modifications by Kumo Inc.
// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//



#pragma once

#include <atomic>
#include <kmhdfs/common/exception_internal.h>
#include <kmhdfs/proto/ipc_connection_context.pb.h>
#include <memory>
#include <kmhdfs/network/buffered_socket_reader.h>
#include <kmhdfs/network/tcp_socket.h>
#include <kmhdfs/rpc/rpc_call.h>
#include <kmhdfs/rpc/rpc_channel_key.h>
#include <kmhdfs/proto/rpc_header.pb.h>
#include <kmhdfs/rpc/rpc_remote_call.h>
#include <kmhdfs/rpc/sasl_client.h>
#include <kmhdfs/common/thread.h>

#include <google/protobuf/message.h>

namespace Hdfs {
    namespace Internal {
        class RpcClient;

        class RpcChannel {
        public:
            /**
             * Destroy a channel
             */
            virtual ~RpcChannel() {
            }

            /**
             * The caller finished the rpc call,
             * this channel may be reused later if immediate is false.
             * @param immediate Do not reuse the channel any more if immediate is true.
             */
            virtual void close(bool immediate) = 0;

            /**
             * Invoke a rpc call.
             * @param call The call is to be invoked.
             * @return The remote call object.
             */
            virtual void invoke(const RpcCall &call) = 0;

            /**
             * Close the channel if it idle expired.
             * @return true if the channel idle expired.
             */
            virtual bool checkIdle() = 0;

            /**
             * Wait for all reference exiting.
             * The channel cannot be reused any more.
             * @pre RpcClient is not running.
             */
            virtual void waitForExit() = 0;

            /**
             * Add reference count to this channel.
             */
            virtual void addRef() = 0;
        };

        /**
         * RpcChannel represent a rpc connect to the server.
         */
        class RpcChannelImpl : public RpcChannel {
        public:
            /**
             * Construct a RpcChannelImpl instance.
             * @param k The key of this channel.
             */
            RpcChannelImpl(const RpcChannelKey &k, RpcClient &c);

            /**
             * Destroy a RpcChannelImpl instance.
             */
            ~RpcChannelImpl();

            /**
             * The caller finished the rpc call,
             * this channel may be reused later if immediate is false.
             * @param immediate Do not reuse the channel any more if immediate is true.
             */
            void close(bool immediate);

            /**
             * Invoke a rpc call.
             * @param call The call is to be invoked.
             * @return The remote call object.
             */
            void invoke(const RpcCall &call);

            /**
             * Close the channel if it idle expired.
             * @return true if the channel idle expired.
             */
            bool checkIdle();

            /**
             * Wait for all reference exiting.
             * The channel cannot be reused any more.
             * @pre RpcClient is not running.
             */
            void waitForExit();

            /**
             * Add reference count to this channel.
             */
            void addRef() {
                ++refs;
            }

        private:
            /**
             * Setup the RPC connection.
             * @pre Already hold write lock.
             */
            void connect();

            /**
             * Cleanup all pending calls.
             * @param reason The reason to cancel the call.
             * @pre Already hold write lock.
             */
            void cleanupPendingCalls(exception_ptr reason);

            /**
             * Send rpc connect protocol header.
             * @throw HdfsNetworkException
             * @throw HdfsTimeout
             */
            void sendConnectionHeader(const RpcAuth &auth);

            /**
             * Send rpc connection protocol content.
             */
            void sendConnectionContent(const RpcAuth &auth);

            /**
             * Build rpc connect context.
             */
            void buildConnectionContext(IpcConnectionContextProto &connectionContext, const RpcAuth &auth);

            /**
             * Send ping packet to server.
             * @throw HdfsNetworkException
             * @throw HdfsTimeout
             * @pre Caller should hold the write lock.
             */
            void sendPing();

            /**
             * Send the call message to rpc server.
             * @param remote The remote call.
             * @pre Already hold write lock.
             */
            void sendRequest(RpcRemoteCallPtr remote);

            /**
             * Issue a rpc call and check response.
             * Catch all recoverable error in this function
             *
             * @param remote The remote call
             */
            exception_ptr invokeInternal(RpcRemoteCallPtr remote);

            /**
             * Check response, block until get one response.
             * @pre Channel already hold read lock.
             */
            void checkOneResponse();

            /**
             * read and handle one response.
             * @pre Channel already hold read lock.
             */
            void readOneResponse(bool writeLock);

            /**
             * Get the call object with given id, and then remove it from pending call list.
             * @param id The id of the call object to be returned.
             * @return The call object with given id.
             * @throw HdfsIOException
             * @pre Channel already locked.
             */
            RpcRemoteCallPtr getPendingCall(int32_t id);

            /**
             * Check if there is data available for reading on socket.
             * @return true if response is available.
             */
            bool getResponse();

            /**
             * wake up one caller to check response.
             * @param id The call id which current caller handled.
             */
            void wakeupOneCaller(int32_t id);

            /**
             * shutdown the RPC connection since error.
             * @param reason The reason to cancel the call
             * @pre Already hold write lock.
             */
            void shutdown(exception_ptr reason);

            const RpcSaslProto_SaslAuth *createSaslClient(
                const ::google::protobuf::RepeatedPtrField<RpcSaslProto_SaslAuth> *auths);

            void sendSaslMessage(RpcSaslProto *msg, ::google::protobuf::Message *resp);

            std::string saslEvaluateToken(RpcSaslProto &response, bool serverIsDone);

            RpcAuth setupSaslConnection();

        private:
            /**
             * Construct a RpcChannelImpl instance for test.
             * @param key The key of this channel.
             * @param sock The socket instance.
             * @param in The BufferedSocketReader instance build on sock.
             * @param client The RpcClient instance.
             */
            RpcChannelImpl(const RpcChannelKey &key, Socket *sock,
                           BufferedSocketReader *in, RpcClient &client);

        private:
            std::atomic<int> refs;
            bool available;
            mutex readMut;
            mutex writeMut;
            RpcChannelKey key;
            RpcClient &client;
            std::shared_ptr<BufferedSocketReader> in;
            std::shared_ptr<SaslClient> saslClient;
            std::shared_ptr<Socket> sock;
            std::atomic<int64_t> lastActivity;
            // ping is a kind of activity, lastActivity will be updated after ping
            std::atomic<int64_t> lastIdle;
            // ping cannot change idle state. If there is still pending calls, lastIdle is always "NOW".
            std::unordered_map<int32_t, RpcRemoteCallPtr> pendingCalls;
        };
    }
}
